# tools/mcp_tools.py
# MCP工具初始化和包装器

import asyncio
import logging
import traceback
from pathlib import Path
from typing import List, Any

from .mcp_client import MCPClientManager
from .mcp_import_manager import mcp_import_manager


class MCPToolsManager:
    """MCP工具管理器，负责初始化和包装MCP工具"""
    
    def __init__(self):
        self.client_manager = MCPClientManager()
        self.tools = []
        self.import_manager = mcp_import_manager
        self.logger = self._setup_logger()
    
    def _setup_logger(self) -> logging.Logger:
        """设置日志记录器"""
        logger = logging.getLogger("mcp_tools_manager")
        logger.setLevel(logging.DEBUG)
        
        if not logger.handlers:
            # 控制台处理器
            console_handler = logging.StreamHandler()
            console_formatter = logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
            )
            console_handler.setFormatter(console_formatter)
            console_handler.setLevel(logging.INFO)
            logger.addHandler(console_handler)
            
            # 文件处理器 - 详细日志
            log_file = Path("mcp_tools_debug.log")
            file_handler = logging.FileHandler(log_file, encoding='utf-8')
            file_formatter = logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(funcName)s:%(lineno)d - %(message)s'
            )
            file_handler.setFormatter(file_formatter)
            file_handler.setLevel(logging.DEBUG)
            logger.addHandler(file_handler)
            
        return logger
    
    async def initialize_tools(self, base_dir: Path) -> List[Any]:
        """初始化MCP工具，支持从配置文件导入"""
        try:
            # 首先尝试从导入的配置初始化服务器
            import_results = await self.import_manager.initialize_enabled_servers()
            
            if import_results:
                print(f"从配置导入的服务器: {list(import_results.keys())}")
                langchain_tools = []
                
                # 为每个导入的服务器创建工具包装器
                for server_id, result in import_results.items():
                    if result.get("success") and "tools" in result:
                        print(f"服务器 {server_id} 初始化成功，工具数量: {len(result['tools'])}")
                        for tool in result["tools"]:
                            wrapper = self._create_imported_tool_wrapper(server_id, tool)
                            langchain_tools.append(wrapper)
                    else:
                        error_msg = result.get("error", "未知错误")
                        print(f"服务器 {server_id} 初始化失败: {error_msg}")
                
                # 不再添加复合工具，让系统使用原有的音乐播放工具
                
                if langchain_tools:
                    print(f"成功创建 {len(langchain_tools)} 个工具包装器")
                    self.tools = langchain_tools
                    return langchain_tools
                else:
                    print("没有成功创建任何工具包装器")
            else:
                print("没有找到任何启用的服务器配置")
            
            # 如果导入的配置没有可用的工具，回退到默认的math服务器
            print("回退到默认的math服务器...")
            return await self._initialize_default_tools(base_dir)
                    
        except Exception as e:
            print(f"MCP工具初始化失败: {e}")
            import traceback
            traceback.print_exc()
            # 回退到默认初始化
            return await self._initialize_default_tools(base_dir)
    
    async def _initialize_default_tools(self, base_dir: Path) -> List[Any]:
        """初始化默认的MCP工具"""
        math_server_path = str(base_dir / "tools" / "math_mcp_server.py")
        if not Path(math_server_path).exists():
            raise FileNotFoundError(f"找不到 MCP 服务器文件: {math_server_path}")
        
        try:
            # 初始化MCP客户端管理器
            success = await self.client_manager.initialize(math_server_path)
            if not success:
                return []
            
            # 获取可用工具
            mcp_tools = await self.client_manager.list_tools()
            print(f"成功加载 MCP 工具: {[tool.name for tool in mcp_tools]}")
            
            # 将MCP工具转换为LangChain工具格式
            from langchain_core.tools import tool
            langchain_tools = []
            
            # 为每个工具创建独立的包装器
            for mcp_tool in mcp_tools:
                wrapper = self._create_tool_wrapper(mcp_tool)
                langchain_tools.append(wrapper)
            
            self.tools = langchain_tools
            return langchain_tools
                    
        except Exception as e:
            print(f"默认MCP工具初始化失败: {e}")
            import traceback
            traceback.print_exc()
            return []
    
    def _create_tool_wrapper(self, mcp_tool: Any) -> Any:
        """为MCP工具创建LangChain包装器"""
        from langchain_core.tools import tool
        
        tool_name = mcp_tool.name
        tool_description = mcp_tool.description
        
        # 通用包装器，使用**kwargs处理所有工具
        @tool
        async def mcp_tool_wrapper(**kwargs):
            """MCP工具包装器 - 通用"""
            print(f"[工具包装器调试] 工具: {tool_name}, 接收到的参数: {kwargs}")
            
            # 处理参数嵌套问题：如果kwargs中有'kwargs'键，则使用其值
            if 'kwargs' in kwargs and len(kwargs) == 1:
                actual_kwargs = kwargs['kwargs']
                print(f"[工具包装器调试] 检测到参数嵌套，使用实际参数: {actual_kwargs}")
                return await self.client_manager.call_tool(tool_name, actual_kwargs)
            else:
                return await self.client_manager.call_tool(tool_name, kwargs)
        
        # 设置工具属性
        mcp_tool_wrapper.name = tool_name
        mcp_tool_wrapper.description = tool_description
        return mcp_tool_wrapper
    
    def _create_imported_tool_wrapper(self, server_id: str, mcp_tool: Any) -> Any:
        """为导入的MCP工具创建LangChain包装器，使用动态参数定义"""
        from langchain_core.tools import StructuredTool
        from pydantic import BaseModel, Field, create_model
        from typing import Optional
        
        tool_name = mcp_tool.name
        tool_description = mcp_tool.description or f"{tool_name} 工具"
        
        self.logger.debug(f"创建工具包装器: {server_id}.{tool_name}")
        print(f"创建工具包装器: {server_id}.{tool_name}")
        
        # 解析工具的输入schema
        input_schema = getattr(mcp_tool, 'inputSchema', None)
        
        # 创建Pydantic模型来定义参数
        if input_schema and isinstance(input_schema, dict):
            properties = input_schema.get('properties', {})
            required = input_schema.get('required', [])
            
            # 动态创建参数字段
            fields = {}
            for prop_name, prop_info in properties.items():
                prop_type = prop_info.get('type', 'string')
                prop_desc = prop_info.get('description', '')
                
                # 映射JSON Schema类型到Python类型
                if prop_type == 'string':
                    py_type = str
                elif prop_type == 'integer':
                    py_type = int
                elif prop_type == 'number':
                    py_type = float
                elif prop_type == 'boolean':
                    py_type = bool
                else:
                    py_type = str
                
                # 如果不是必需参数，使用Optional
                if prop_name not in required:
                    fields[prop_name] = (Optional[py_type], Field(default=None, description=prop_desc))
                else:
                    fields[prop_name] = (py_type, Field(description=prop_desc))
            
            # 创建动态模型
            InputModel = create_model(
                f'{tool_name}Input',
                **fields
            )
        else:
            # 如果没有schema，创建一个通用的模型
            InputModel = create_model(
                f'{tool_name}Input',
                **{}  # 空参数
            )
        
        # 创建异步调用函数
        async def tool_func(**kwargs) -> str:
            """工具调用函数"""
            self.logger.debug(f"工具调用: 服务器={server_id}, 工具={tool_name}, 参数={kwargs}")
            print(f"🔧 [工具调用] {server_id}.{tool_name}")
            print(f"📝 [参数] {kwargs}")
            
            try:
                # 处理参数
                processed_kwargs = self._process_tool_parameters(server_id, tool_name, kwargs)
                self.logger.debug(f"处理后的参数: {processed_kwargs}")
                
                # 调用工具
                result = await self.import_manager.call_tool(server_id, tool_name, **processed_kwargs)
                self.logger.debug(f"工具调用结果: {result}")
                print(f"✅ [结果] {result[:200] if len(str(result)) > 200 else result}")
                return result
            except Exception as e:
                error_msg = f"工具调用失败: {str(e)}"
                self.logger.error(error_msg)
                self.logger.error(f"异常详情: {traceback.format_exc()}")
                print(f"❌ [错误] {error_msg}")
                return error_msg
        
        # 使用StructuredTool创建工具
        structured_tool = StructuredTool(
            name=f"{server_id}_{tool_name}",
            description=f"[{server_id}] {tool_description}",
            coroutine=tool_func,
            args_schema=InputModel
        )
        
        return structured_tool
    
    
    def _process_tool_parameters(self, server_id: str, tool_name: str, kwargs: dict) -> dict:
        """处理工具参数，进行必要的参数名称映射"""
        self.logger.debug(f"处理参数: {server_id}.{tool_name}, 输入: {kwargs}")
        processed_kwargs = kwargs.copy()
        
        # 参数映射规则
        param_mappings = {
            ("bingcn", "bing_search"): {"q": "query"},
            ("weather-server", "search_netease_music"): {"keywords": "keyword"}
        }
        
        # 应用参数映射
        mapping_key = (server_id, tool_name)
        if mapping_key in param_mappings:
            for old_key, new_key in param_mappings[mapping_key].items():
                if old_key in processed_kwargs:
                    processed_kwargs[new_key] = processed_kwargs.pop(old_key)
                    self.logger.debug(f"参数映射: {old_key} -> {new_key}")
        
        # 天气工具默认参数
        if server_id == "weather-server" and tool_name == "weather_query":
            if "city" in processed_kwargs and "days" not in processed_kwargs:
                processed_kwargs["days"] = 1
                self.logger.debug("设置天气工具默认days=1")
        
        self.logger.debug(f"处理后参数: {processed_kwargs}")
        return processed_kwargs
    
    async def close(self):
        """关闭MCP连接 - 添加异常保护"""
        try:
            await asyncio.wait_for(self.client_manager.close(), timeout=2.0)
        except Exception as e:
            self.logger.warning(f"关闭客户端管理器失败: {e}")
        
        try:
            await asyncio.wait_for(self.import_manager.close_all_connections(), timeout=5.0)
        except Exception as e:
            self.logger.warning(f"关闭导入管理器失败: {e}")


# 全局MCP工具管理器实例
mcp_tools_manager = MCPToolsManager()


async def initialize_mcp_tools(base_dir: Path) -> List[Any]:
    """初始化MCP工具的便捷函数"""
    return await mcp_tools_manager.initialize_tools(base_dir)


async def close_mcp_tools():
    """关闭MCP工具的便捷函数"""
    await mcp_tools_manager.close()
